#include "octreequantizer.h"
#include <QException>

namespace WingGif {

const int Octree::mask[] = {0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01};

Octree::OctreeNode::OctreeNode(int level, int colorBits, Octree *octree) {
  // Construct the new node.
  _leaf = level == colorBits;

  _red = _green = _blue = 0;
  _pixelCount = 0;

  // If a leaf, increment the leaf count.
  if (_leaf) {
    octree->_leaves++;
    _nextReducible = nullptr;
    _children.clear();
  } else {
    // Otherwise add this to the reducible nodes.
    auto reducibleNodes = octree->reducibleNodes();

    _nextReducible = reducibleNodes[level];
    reducibleNodes[level] = this;
    _children.resize(8);
  }
  setParent(octree);
}

void Octree::OctreeNode::addColor(QColor pixel, int colorBits, int level,
                                  Octree *octree) {
  // Update the color information if this is a leaf
  if (_leaf) {
    increment(pixel);

    // Setup the previous node.
    octree->trackPrevious(this);
  } else {
    // Go to the next level down in the tree.
    auto shift = 7 - level;
    auto index = ((pixel.red() & mask[level]) >> (shift - 2)) |
                 ((pixel.green() & mask[level]) >> (shift - 1)) |
                 ((pixel.blue() & mask[level]) >> (shift));

    auto child = _children[index];

    if (!child) {
      // Create a new child node & store in the array.
      child = new OctreeNode(level + 1, colorBits, octree);
      _children[index] = child;
    }

    // Add the color to the child node.
    child->addColor(pixel, colorBits, level + 1, octree);
  }
}

Octree::OctreeNode *Octree::OctreeNode::nextReducible() {
  return _nextReducible;
}

QVector<Octree::OctreeNode *> &Octree::OctreeNode::children() {
  return _children;
}

int Octree::OctreeNode::reduce() {
  _red = _green = _blue = 0;
  auto children = 0;

  // Loop through all children and add their information to this node.
  for (auto index = 0; index < 8; index++) {
    if (!_children[index])
      continue;

    _red += _children[index]->_red;
    _green += _children[index]->_green;
    _blue += _children[index]->_blue;
    _pixelCount += _children[index]->_pixelCount;
    ++children;

    _children[index]->deleteLater();
    _children[index] = nullptr;
  }

  // Now change this to a leaf node.
  _leaf = true;

  // Return the number of nodes to decrement the leaf count by.
  return children - 1;
}

void Octree::OctreeNode::constructPalette(QVector<QColor> &palette,
                                          int &paletteIndex) {
  if (_leaf) {
    // Consume the next palette index.
    _paletteIndex = paletteIndex++;

    // And set the color of the palette entry.
    palette.append(QColor::fromRgb(_red / _pixelCount, _green / _pixelCount,
                                   _blue / _pixelCount));
  } else {
    // Loop through children looking for leaves.
    for (auto index = 0; index < 8; index++) {
      if (_children[index])
        _children[index]->constructPalette(palette, paletteIndex);
    }
  }
}

int Octree::OctreeNode::getPaletteIndex(QColor pixel, int level) {
  auto paletteIndex = _paletteIndex;

  if (_leaf)
    return paletteIndex;

  auto shift = 7 - level;
  auto index = ((pixel.red() & mask[level]) >> (shift - 2)) |
               ((pixel.green() & mask[level]) >> (shift - 1)) |
               ((pixel.blue() & mask[level]) >> (shift));

  if (_children[index])
    paletteIndex = _children[index]->getPaletteIndex(pixel, level + 1);
  else
    throw "Not expected!";

  return paletteIndex;
}

void Octree::OctreeNode::increment(QColor pixel) {
  _pixelCount++;
  _red += pixel.red();
  _green += pixel.green();
  _blue += pixel.blue();
}

Octree::Octree(int maxColorBits) {
  _maxColorBits = maxColorBits;
  _leaves = 0;
  _reducibleNodes.resize(9);
  _root = new OctreeNode(0, _maxColorBits, this);
  _previousColor = Qt::transparent;
  _previousNode = nullptr;
}

QVector<Octree::OctreeNode *> &Octree::reducibleNodes() {
  return _reducibleNodes;
}

void Octree::reduce() {
  int index;

  // Find the deepest level containing at least one reducible node
  for (index = _maxColorBits - 1; index > 0 && !_reducibleNodes[index]; index--)
    ;

  // Reduce the node most recently added to the list at level 'index'
  auto node = _reducibleNodes[index];
  _reducibleNodes[index] = node->nextReducible();

  // Decrement the leaf count after reducing the node
  _leaves -= node->reduce();

  // And just in case I've reduced the last color to be added, and the next
  // color to be added is the same, invalidate the previousNode...
  _previousNode->deleteLater();
  _previousNode = nullptr;
}

void Octree::trackPrevious(Octree::OctreeNode *node) { _previousNode = node; }

QVector<QColor> Octree::palletize(int colorCount) {
  while (_leaves > colorCount)
    reduce();

  // Now palettize the nodes.
  auto palette = QVector<QColor>(_leaves);

  auto paletteIndex = 0;

  _root->constructPalette(palette, paletteIndex);

  // And return the palette.
  return palette;
}

int Octree::getPaletteIndex(QColor pixel) {
  return _root->getPaletteIndex(pixel, 0);
}

void Octree::addColor(QColor pixel) {
  // Check if this request is for the same color as the last
  if (_previousColor == pixel) {
    // If so, check if I have a previous node setup. This will only occur if the
    // first color in the image happens to be black, with an alpha component of
    // zero.
    if (!_previousNode) {
      _previousColor = pixel;
      _root->addColor(pixel, _maxColorBits, 0, this);
    } else
      // Just update the previous node
      _previousNode->increment(pixel);
  } else {
    _previousColor = pixel;
    _root->addColor(pixel, _maxColorBits, 0, this);
  }
}

OctreeQuantizer::OctreeQuantizer(int maxColorBits) : Quantizer(false, this) {
  if (maxColorBits < 1 || maxColorBits > 8)
    throw "This should be between 1 and 8";

  // Construct the octree.
  _octree = new Octree(maxColorBits);
}

OctreeQuantizer::~OctreeQuantizer() { delete _octree; }

void OctreeQuantizer::initialQuantizePixel(QColor pixel) {
  if (pixel.red() == 0)
    return;

  // Add the color to the octree.
  _octree->addColor(pixel);
}

uchar OctreeQuantizer::quantizePixel(QColor pixel) {
  return uchar(_octree->getPaletteIndex(pixel));
}

QVector<QColor> OctreeQuantizer::buildPalette() {
  auto _maxColors = maxColors();
  auto _transparentColor = transparentColor();
  auto maxColorsWithTransparency =
      _transparentColor.isValid() ? _maxColors - 1 : _maxColors;
  setMaxColorsWithTransparency(maxColorsWithTransparency);

  // First off convert the octree to _maxColors colors
  auto palette = _octree->palletize(maxColorsWithTransparency);

  // TODO: Since the color table changes in size by ^2 (64, 128, 256), if there
  // is still space in the color table, there's no need for the (-1). Check

  // Add the transparent color to the last position.
  if (_transparentColor.isValid())
    palette.append(QColor::fromRgba(qRgba(_transparentColor.red(),
                                          _transparentColor.green(),
                                          _transparentColor.blue(), 0)));
  // I need to set a color that is not
  // being used in the gif.

  // Just convert the array to a list.
  return palette;
}

} // namespace WingGif
